#include <kernel.h>
#include <malloc.h>
#include <string.h>
#include <csp_debug.h>
#include <csp_macros.h>
#include <csp_protocol_msg.h>
#include <kdma.h>
#include <syscall.h>
#include <process_control.h>
#include <profiling.h>


/* Save area for kernels registers r0, r1 and r2 */
qword kernel_pointers[2];
qword kernel_stackpointer;

/* Used for calculating kernel stack starting point. */
register qword kernel_sp asm("$sp");
/* kernel size */
extern unsigned int _end;

/* Kernel initial setup  */
spe_args_t setup __attribute__ ((aligned(128)));

/* Addresses of current pcb's in EA */
ea_addr_t pcb_ea_ptrs[MAX_SPU_SLOTS] __attribute__ ((aligned(128))); 

//process_registers_t *current;
spu_pcb_t *current;

/* Process control blocks */
spu_pcb_t procs[MAX_SPU_SLOTS];

/* Use this space for swapping pcbs*/
pcb_t swap_pcb_space[MAX_SPU_SLOTS];

/* This slot is executing */
uint32_t current_slot, counter;

/**/
char completion_flag[2];


void main_loop(int msg, int val);



void write_process_hashes(uint32_t slot, uint64_t* hashes)
{
	int size = (int) hashes[0];

	large_put_blocking(
			(ls_addr_t) hashes,
			setup.offload_area_ea,
			CEIL_16( (sizeof(uint64_t) * (size+1)) ),
			slot);
}

void write_completion_flag(pcb_t *pcb)
{
	little_put_barrier(
			(ls_addr_t) completion_flag,
			pcb->dma_completion_flag_ea_addr,
			sizeof(char) * 2,
			pcb->slot);
}

void update_pcb_ptrs(uint32_t slot)
{
	//	uint32_t floor_even = slot - (slot % 2);

	little_get_blocking(
			setup.process_data_EA,
			(ls_addr_t) &pcb_ea_ptrs[0], 
			sizeof(ea_addr_t)*MAX_SPU_SLOTS,
			slot);
}

void update_pcb(spu_pcb_t *spu_pcb)
{
	ls_addr_t sp_addr, avail_stack;

	/* Calculate how much space we need for stack. */
	sp_addr = si_to_int(  spu_pcb->d_regs[1] ) - NCALLEESAVES*sizeof(qword);
	avail_stack = ( (int32_t*)  &spu_pcb->d_regs[1])[1] /*-
														  (spu_pcb->process_data.vaddr + spu_pcb->process_data.exec_size + spu_pcb->process_data.bss_size) 
														  */;	
		spu_pcb->process_data.stack_size = 
		mfc_ceil128( spu_pcb->process_data.stack_start - sp_addr );

	spu_pcb->process_data.heap_size = ( sp_addr - avail_stack ) - 
		(spu_pcb->process_data.vaddr + 
		 spu_pcb->process_data.exec_size + 
		 spu_pcb->process_data.bss_size);
}


void write_pcb(uint32_t slot)
{
	spu_pcb_t *pcb = &procs[slot];

	ASSERT( 0 == pcb->process_data.dma_completion_flag[0]);

	little_put(
			(ls_addr_t) &pcb->process_data, 
			pcb->process_data.pcb_ea, 
			sizeof(pcb_t), 
			slot
			);
	write_completion_flag(&pcb->process_data);
}

void read_pcb(uint32_t slot) 
{

	little_get_blocking(
			pcb_ea_ptrs[slot], 
			(ls_addr_t) &procs[slot].process_data, 
			sizeof(pcb_t),
			slot);

	procs[slot].has_process = 1;
}



void read_process(spu_pcb_t *procs, uint32_t slot)
{
	spu_pcb_t *pcb = &procs[slot];


	ls_addr_t ls_save_area_regs, ls_save_area_stack, ls_save_area_bss; 

	ea_addr_t ea_save_area_size, ea_save_area_regs, ea_save_area_stack, 
			  ea_save_area_bss;

	ls_size_t save_area_size_regs, save_area_size_stack, save_area_size_bss;

	ea_save_area_size = SAVE_PROC_AREA_SIZE( pcb->process_data );
	ea_save_area_regs = SAVE_REG_AREA( pcb->process_data );
	ea_save_area_stack= SAVE_STACK_AREA( pcb->process_data );
	ea_save_area_bss  = SAVE_BSS_AREA( pcb->process_data );

	save_area_size_regs = VOLATILEREGISTERS*sizeof(qword);
	save_area_size_stack= pcb->process_data.stack_size;
	save_area_size_bss  = CEIL_16(pcb->process_data.bss_size);

	ls_save_area_regs  = (ls_addr_t) &pcb->d_regs;
	ls_save_area_stack = pcb->process_data.stack_start - pcb->process_data.stack_size;
	ls_save_area_bss   = pcb->process_data.vaddr + pcb->process_data.exec_size;


	if( pcb->program != pcb->process_data.exec ) {

		large_get(
				pcb->process_data.exec, 
				pcb->process_data.vaddr, 
				CEIL_16(pcb->process_data.exec_size), 
				slot);
		pcb->program = pcb->process_data.exec;
	}


	if(pcb->process_data.state != STATE_INIT){


		DEBUG_PRINTF("READ(slot=%i, pid=%i): size=%llx;"
				"regs_ea=%llx, size=%x, regs_ls=%p;" 
				"stack=%llx, size=%x, stack_ls=%p;"
				"bss=%llx, size=%x, bss_ls=%p\n", 
				slot, pcb->process_data.pid, ea_save_area_size,  
				ea_save_area_regs, save_area_size_regs, ls_save_area_regs,
				ea_save_area_stack,  save_area_size_stack, ls_save_area_stack,
				ea_save_area_bss, save_area_size_bss , ls_save_area_bss);




		/* registers */
		little_get(
				ea_save_area_regs, 
				ls_save_area_regs, 
				save_area_size_regs, 
				slot);

		/* stack */
		large_get(
				ea_save_area_stack, 
				ls_save_area_stack, 
				save_area_size_stack,
				slot);

		/* bss + heap */
		if( pcb->process_data.bss_size || pcb->process_data.heap_size ){

			large_get(
					ea_save_area_bss, 
					ls_save_area_bss, 
					save_area_size_bss + pcb->process_data.heap_size ,
					slot
					);
		}
	}else{
		/* Clear bss before init. */
		memset( (void*) 
				pcb->process_data.vaddr + pcb->process_data.exec_size,
				0, pcb->process_data.bss_size);
	}
}

void write_process(spu_pcb_t *procs, uint32_t slot)
{
	//uint32_t stack_top; 
	spu_pcb_t *pcb = &procs[slot];

	ea_addr_t ea_save_area_size, ea_save_area_regs, 
			  ea_save_area_stack, ea_save_area_bss; 

	ls_size_t save_area_size_regs, save_area_size_stack, 
			  save_area_size_bss  ;

	ls_addr_t ls_save_area_regs, ls_save_area_stack, ls_save_area_bss; 


	ea_save_area_size = SAVE_PROC_AREA_SIZE( pcb->process_data );
	ea_save_area_regs = SAVE_REG_AREA( pcb->process_data );
	ea_save_area_stack= SAVE_STACK_AREA( pcb->process_data );
	ea_save_area_bss  = SAVE_BSS_AREA( pcb->process_data );

	save_area_size_regs = VOLATILEREGISTERS*sizeof(qword);
	save_area_size_stack= pcb->process_data.stack_size;
	save_area_size_bss  = CEIL_16(pcb->process_data.bss_size);

	ls_save_area_regs  = (ls_addr_t) &pcb->d_regs;
	ls_save_area_stack = pcb->process_data.stack_start - pcb->process_data.stack_size;
	ls_save_area_bss   = pcb->process_data.vaddr + pcb->process_data.exec_size;


	DEBUG_PRINTF("WRITE(slot=%i, pid=%i): size=%llx;"
			"regs_ea=%llx, size=%x, regs_ls=%p;"
			"stack=%llx, size=%x, stack_ls=%p;"
			"bss=%llx, size=%x, bss_ls=%p\n", 
			slot, pcb->process_data.pid, ea_save_area_size,  
			ea_save_area_regs, save_area_size_regs, ls_save_area_regs,
			ea_save_area_stack,  save_area_size_stack, ls_save_area_stack,
			ea_save_area_bss, save_area_size_bss , ls_save_area_bss);



	/* registers */
	little_put(
			ls_save_area_regs, 
			ea_save_area_regs, 
			save_area_size_regs, 
			slot);

	/* stack */
	large_put( ls_save_area_stack , ea_save_area_stack , save_area_size_stack , 
			slot);

	/* bss + heap */
	if( pcb->process_data.bss_size || pcb->process_data.heap_size ){ // Only if we have bss

		large_put( 
				ls_save_area_bss, 
				ea_save_area_bss, 
				save_area_size_bss  + pcb->process_data.heap_size, 
				slot);
	}
}



void get_process(uint32_t slot)
{

#ifdef CLEAR_SPE_MEM
	procs[slot].program = 0;
	memset( (void*) procs[slot].process_data.vaddr, 
			0, 
			procs[slot].process_data.stack_start - procs[slot].process_data.vaddr);
	memset(&procs[slot], 0, sizeof(spu_pcb_t));
#endif
	read_pcb(slot);
	//read_program(slot);
	read_process(procs, slot);
}

int32_t handle_channel_reply(csp_channel_t *channel)
{

	csp_channel_t *chan;
	int32_t retval = -1;

	uint32_t msg = spu_read_in_mbox();

	uint32_t cmd = MSG_CMD(msg);
	uint32_t option = MSG_OPT(msg);


	switch( cmd ){

		case COMM_POISON:

			/*We have tried a channel-op on a poisoned channel. */ 				
			/* Tell process channel is poisoned. */
			current->process_data.op_ret_val = CHANNEL_POISON;

			__resume(current->process_data.op_ret_val);
			break;
		case COMM_CHANN_YES:
			/* Get corr. process' pcb */
			little_get_blocking(setup.pcb_head + (sizeof(pcb_t)*option), 
					(ls_addr_t) &swap_pcb_space[current->process_data.slot], 
					sizeof(pcb_t), current->process_data.slot );

			chan = &channel[swap_pcb_space[current->process_data.slot].misc_val];

			/* this can be done without branch by a function pointer */
			if( CHANNEL_READ == chan->op ){ 
				large_get_blocking(
						swap_pcb_space[current->process_data.slot].channel_data,
						chan->addr,
						chan->nbytes,
						current->process_data.slot);
			}else{ /* write */
				large_put_blocking( 
						chan->addr,
						swap_pcb_space[current->process_data.slot].channel_data,
						chan->nbytes,
						current->process_data.slot);	
			}

			/* Return value when process is resumed. */
			current->process_data.op_ret_val = 
				swap_pcb_space[current->process_data.slot].misc_val;


			ASSERT( 0 == swap_pcb_space[current->process_data.slot].dma_completion_flag[0]);
			ASSERT( STATE_YIELD == swap_pcb_space[current->process_data.slot].state );

			ASSERT( STATE_YIELD == swap_pcb_space[current->process_data.slot].state );

			write_completion_flag(&swap_pcb_space[current->process_data.slot]);

			/* Execute other process to hide latency*/
			retval = COMM_EXECUTE;

			break;	
		case COMM_CHANN_NO_YIELD:

			current->process_data.state = STATE_CHAN_WAIT;
			retval = PROCESS_SWAP;
			break;
		default:
			printf("SPU #%d: handle_channel_reply: protocol error=0x%x.\n", setup.spe_id, msg);
			exit(-1);
	}

	return retval;
}

int32_t do_alting(int val)
{ 
	csp_guard_list_t *gl = (csp_guard_list_t*) val;

	current->process_data.state = STATE_YIELD;

	little_put_blocking((ls_addr_t) gl->head, setup.offload_area_ea, 
			sizeof(csp_channel_t)*gl->size, current->process_data.slot);

	/* Hack in msg: slot is a number. */
	write_mbox_blocking( COMM_ALT | (current->process_data.slot << 12) | gl->size );

	return handle_channel_reply(gl->head);
}

int32_t do_channel_communication(ls_addr_t channel)
{
	current->process_data.state = STATE_YIELD;
	little_put_blocking(channel, setup.offload_area_ea, sizeof(csp_channel_t), current->process_data.slot);

	write_mbox_blocking(COMM_CHANN_TALK | current->process_data.slot);

	return handle_channel_reply((csp_channel_t*) channel);
}





int32_t do_parallel(ls_addr_t* package){
	int payload_size;
	char** mnemonics;
	uint64_t *mnemonics_hashes;

	//_PRINT_EXIT("NOT HERE!\n");
	DEBUG_PRINTF("SPU #%d: COMM_INIT_PROCS "
			"(pid=%i, state=%x, slot=%i)\n", 
			setup.spe_id,
			current->process_data.pid,
			current->process_data.state,
			current->process_data.slot);

	//package = (ls_addr_t*) val;

	payload_size = (int) package[0];
	mnemonics = (char**) package[1];
	mnemonics_hashes = (uint64_t*) package[2];

	write_process_hashes(current->process_data.slot, mnemonics_hashes);

	current->process_data.state = STATE_WAIT;
	write_mbox_blocking(COMM_INIT_PROCS | current->process_data.slot);

	update_pcb(current);
	write_process(procs, current->process_data.slot);
	write_pcb(current->process_data.slot);
	write_mbox_blocking(COMM_PUT_PROC | current->process_data.state | current->process_data.slot);

	current->has_process = 0;

	return GET_CMD;
}


/************************ SYSCALLS ON USER STACK *************************/
/* Syscalls done on user-stack must never end save of current process.   */
void csp_poison_channel(ls_addr_t channel){
	printf("poison spe #%i\n", setup.spe_id);
	little_put_blocking(channel, setup.offload_area_ea, sizeof(csp_channel_t), current->process_data.slot);
	write_mbox_blocking(COMM_POISON | current->process_data.slot);
	//main_loop(GET_CMD, 0);
}

ea_addr_t csp_shared_address(){
	return setup.shared_addr;
}

int csp_spe_id(){
	return setup.spe_id;
}


void csp_exit(int ret   __attribute__ ((unused))){
	printf("WE SHOULD NOT GO HERE!!\n");
///*
//	current->process_data.state = STATE_EXIT;
//	if( !current->process_data.ret_val ){
//		current->process_data.ret_val = ret;
//	}
//	DEBUG_PRINTF( "SPU #%d: Exit process, pid=%i, return value=%i, superior return value=%i.\n", 
//			setup.spe_id, current->process_data.pid, val, current->process_data.ret_val );
//	write_pcb(current->process_data.slot);
//	write_mbox_blocking(COMM_PUT_PROC | current->process_data.state | current->process_data.slot);
//	current->has_process = 0;
//	main_loop(GET_CMD, 0); 
//*/
}


size_t csp_code_size()
{
	return current->process_data.exec_size;
}

void csp_exit_network()
{
	profiling_print(setup.spe_id);
	write_mbox_blocking(COMM_DESTROY);
	exit(0);
	//main_loop(GET_CMD, 0);
}

uint32_t csp_pid()
{
	return current->process_data.pid;
}


syscall_pointers_t syscalls = {
	csp_pid,
	csp_exit_network,
	csp_code_size,
	csp_exit,
	csp_spe_id,
	csp_shared_address,
	csp_poison_channel
};
/******************** END SYSCALLS ON USER STACK *************************/


void main_loop(int msg, int val) 
{
	int t = 20; // Safety bell


	DEBUG_PRINTF("SPU #%d: main_loop\n", setup.spe_id);
	while(t--){
		uint32_t cmd = MSG_CMD(msg);
		uint32_t option = MSG_OPT(msg);


		switch( cmd ) {
			case COMM_CHANN_TALK:
				msg = do_channel_communication(val);
				break;
			case PROCESS_SWAP:
				update_pcb(current);
				write_mbox_blocking(COMM_PUT_PROC | 
						current->process_data.state | 
						current->process_data.slot);
				write_process(procs, current->process_data.slot);
				write_pcb(current->process_data.slot);
				current->has_process = 0;

				msg = GET_CMD;
				break; 
			case COMM_INIT_PROCS:
				msg = do_parallel( (ls_addr_t*) val );
				break; 
			case ALT:
				msg = do_alting(val);
				break;

			case COMM_TAKE_PROC:

				DEBUG_PRINTF("SPU #%d: COMM_TAKE_PROC, option=%i\n", setup.spe_id, option);

				//StartTimer(profile_p);
				update_pcb_ptrs(option);
				get_process(option);
				//StopTimer(profile_p);
				//profile_total += profile_p;

				msg = GET_CMD;
				break;
			case GET_CMD :
				DEBUG_PRINTF("SPU #%d: GET_CMD: Waiting for cmd from PPE..\n", setup.spe_id);
				msg = spu_read_in_mbox();
				break;
			case YIELD : 
				DEBUG_PRINTF( "SPU #%d: Yield process, pid=%i.\n", setup.spe_id, current->process_data.pid );
				current->process_data.state = STATE_YIELD;
				msg = PROCESS_SWAP;
				break;

			case COMM_EXECUTE:
				DEBUG_PRINTF("SPU #%d: COMM_EXECUTE.\n", setup.spe_id);
				/* continue in round-robin fashion until we have tried all slots. */
				int rr = MAX_SPU_SLOTS;
				while(rr--){
					current_slot = ++counter % MAX_SPU_SLOTS;	

					ASSERT(STATE_INIT == current->process_data.state || STATE_YIELD == current->process_data.state);
					if(procs[current_slot].has_process){
						current = &procs[current_slot];


						profiling_dma_start();
						/* Double buffering: make sure data is written before we start executing! */
						mfc_write_tag_mask(1<<current_slot);
						mfc_read_tag_status_all();
						profiling_dma_stop();


						if(STATE_INIT == current->process_data.state){
							DEBUG_PRINTF("SPU #%d: Executing slot=%i, state=INIT.\n", setup.spe_id, current_slot);
							run_process(current, &kernel_stackpointer, &syscalls);
						}else if(STATE_YIELD == current->process_data.state){
							DEBUG_PRINTF("SPU #%d: Executing slot=%i, state=YIELD.\n", setup.spe_id, current_slot);
							__resume(current->process_data.op_ret_val);
						}
					}	
				}

				// We should never go here:
				printf("Protocol error: got COMM_EXECUTE message, but had no process to execute.\n");
				exit(-1);
				break;
			case EXIT:
				current->process_data.state = STATE_EXIT;
				if( !current->process_data.ret_val ){
					current->process_data.ret_val = val;
				}
				DEBUG_PRINTF( "SPU #%d: Exit process, pid=%i, return value=%i, superior return value=%i.\n", 
						setup.spe_id, current->process_data.pid, val, current->process_data.ret_val );
				write_pcb(current->process_data.slot);
				write_mbox_blocking(COMM_PUT_PROC | current->process_data.state | current->process_data.slot);
				current->has_process = 0;
				msg = GET_CMD;
				break;
			case COMM_KILL:
				//printf("SPU #%d: got kill msg.\n", setup.spe_id);
				write_mbox_blocking(COMM_KILL);
				profiling_print(setup.spe_id);
				exit(0);
			default  : 
				printf("SPU #%d: main_loop: protocol error=0x%x.\n", setup.spe_id, msg);
				exit(-1);
				break;
		}	
	}
	exit(-1);
} 


int main(uint64_t _ctx_id  __attribute__ ((unused)), uint64_t _spe_args)
{

	/* Get arguments */
	mfc_get(&setup, _spe_args, sizeof(spe_args_t), TAGID_MAIN, 0, 0);
	mfc_write_tag_mask(1<<TAGID_MAIN);
	mfc_read_tag_status_all();

	completion_flag[0] = PROCESS_DMA_COMPLETION;
	completion_flag[1] = PROCESS_DMA_COMPLETION;

	/* Remember where we can find info on proceses. */
	current = procs;
	current_slot = 0;
	counter = 0;

	/* Profiling */
	profiling_setup();

	/* When returning to kernel-mode: use stack from here. */
	kernel_stackpointer = kernel_sp;
	main_loop(GET_CMD, 0);

	printf("SPU #%i: We should never go here.\n", setup.spe_id);
	return -1;

}
